"""
WebSocket连接自动扩缩容
======================

实现智能的连接扩缩容机制：
1. 负载监控
2. 自动扩容/缩容决策
3. 连接迁移
4. 性能优化
"""

import asyncio
import logging
from typing import Dict, Any, Optional, List
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from enum import Enum

logger = logging.getLogger(__name__)


class ScalingAction(Enum):
    """扩缩容动作"""
    SCALE_UP = "scale_up"      # 扩容
    SCALE_DOWN = "scale_down"  # 缩容
    MAINTAIN = "maintain"       # 保持


class LoadLevel(Enum):
    """负载级别"""
    LOW = "low"          # 低负载 < 30%
    NORMAL = "normal"    # 正常负载 30-70%
    HIGH = "high"        # 高负载 70-90%
    CRITICAL = "critical" # 临界负载 > 90%


@dataclass
class LoadMetrics:
    """负载指标"""
    connection_count: int
    max_connections: int
    cpu_usage_percent: float
    memory_usage_percent: float
    message_rate: float  # 每秒消息数
    error_rate: float    # 错误率
    avg_latency_ms: float
    timestamp: datetime = field(default_factory=datetime.now)
    
    def get_connection_utilization(self) -> float:
        """获取连接使用率"""
        if self.max_connections == 0:
            return 0.0
        return self.connection_count / self.max_connections
    
    def get_load_level(self) -> LoadLevel:
        """获取负载级别"""
        utilization = self.get_connection_utilization()
        
        if utilization < 0.3:
            return LoadLevel.LOW
        elif utilization < 0.7:
            return LoadLevel.NORMAL
        elif utilization < 0.9:
            return LoadLevel.HIGH
        else:
            return LoadLevel.CRITICAL
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "connection_count": self.connection_count,
            "max_connections": self.max_connections,
            "connection_utilization": round(self.get_connection_utilization() * 100, 2),
            "cpu_usage_percent": round(self.cpu_usage_percent, 2),
            "memory_usage_percent": round(self.memory_usage_percent, 2),
            "message_rate": round(self.message_rate, 2),
            "error_rate": round(self.error_rate, 4),
            "avg_latency_ms": round(self.avg_latency_ms, 2),
            "load_level": self.get_load_level().value,
            "timestamp": self.timestamp.isoformat()
        }


@dataclass
class ScalingDecision:
    """扩缩容决策"""
    action: ScalingAction
    reason: str
    target_connections: int
    current_connections: int
    metrics: LoadMetrics
    timestamp: datetime = field(default_factory=datetime.now)
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            "action": self.action.value,
            "reason": self.reason,
            "target_connections": self.target_connections,
            "current_connections": self.current_connections,
            "metrics": self.metrics.to_dict(),
            "timestamp": self.timestamp.isoformat()
        }


class AutoScaler:
    """
    自动扩缩容器
    
    功能：
    - 负载监控
    - 智能决策
    - 自动扩缩容
    - 连接迁移
    """
    
    def __init__(
        self,
        min_connections: int = 100,
        max_connections: int = 10000,
        target_utilization: float = 0.7,
        scale_up_threshold: float = 0.8,
        scale_down_threshold: float = 0.3,
        check_interval_seconds: int = 30,
        cooldown_seconds: int = 300
    ):
        """
        初始化自动扩缩容器
        
        Args:
            min_connections: 最小连接数
            max_connections: 最大连接数
            target_utilization: 目标使用率
            scale_up_threshold: 扩容阈值
            scale_down_threshold: 缩容阈值
            check_interval_seconds: 检查间隔（秒）
            cooldown_seconds: 冷却时间（秒）
        """
        self.min_connections = min_connections
        self.max_connections = max_connections
        self.target_utilization = target_utilization
        self.scale_up_threshold = scale_up_threshold
        self.scale_down_threshold = scale_down_threshold
        self.check_interval_seconds = check_interval_seconds
        self.cooldown_seconds = cooldown_seconds
        
        # 当前容量
        self.current_capacity = min_connections
        
        # 最近的负载指标
        self.recent_metrics: List[LoadMetrics] = []
        
        # 扩缩容历史
        self.scaling_history: List[ScalingDecision] = []
        
        # 最后一次扩缩容时间
        self.last_scaling_time: Optional[datetime] = None
        
        # 后台任务
        self.autoscale_task: Optional[asyncio.Task] = None
        self.is_running = False
        
        # 统计
        self.scale_up_count = 0
        self.scale_down_count = 0
        
        logger.info(
            f"AutoScaler initialized: capacity={min_connections}-{max_connections}, "
            f"target={target_utilization*100}%"
        )
    
    async def start(self):
        """启动自动扩缩容"""
        if self.is_running:
            logger.warning("AutoScaler already running")
            return
        
        self.is_running = True
        self.autoscale_task = asyncio.create_task(self._autoscale_loop())
        
        logger.info("AutoScaler started")
    
    async def stop(self):
        """停止自动扩缩容"""
        self.is_running = False
        
        if self.autoscale_task:
            self.autoscale_task.cancel()
            try:
                await self.autoscale_task
            except asyncio.CancelledError:
                pass
        
        logger.info("AutoScaler stopped")
    
    async def _autoscale_loop(self):
        """自动扩缩容循环"""
        while self.is_running:
            try:
                await asyncio.sleep(self.check_interval_seconds)
                await self._check_and_scale()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Autoscale loop error: {e}", exc_info=True)
    
    async def _check_and_scale(self):
        """检查并执行扩缩容"""
        # 获取当前负载
        metrics = await self._collect_metrics()
        
        if not metrics:
            return
        
        # 记录指标
        self.recent_metrics.append(metrics)
        if len(self.recent_metrics) > 100:
            self.recent_metrics = self.recent_metrics[-100:]
        
        # 检查是否在冷却期
        if self._is_in_cooldown():
            logger.debug("In cooldown period, skipping scaling check")
            return
        
        # 做出扩缩容决策
        decision = self._make_scaling_decision(metrics)
        
        # 执行决策
        if decision.action != ScalingAction.MAINTAIN:
            await self._execute_scaling(decision)
    
    async def _collect_metrics(self) -> Optional[LoadMetrics]:
        """
        收集负载指标
        
        Returns:
            负载指标
        """
        # TODO: 实际实现应该从连接池和系统中收集指标
        # 这里简化为模拟数据
        return LoadMetrics(
            connection_count=500,
            max_connections=self.current_capacity,
            cpu_usage_percent=45.0,
            memory_usage_percent=60.0,
            message_rate=1000.0,
            error_rate=0.001,
            avg_latency_ms=100.0
        )
    
    def _is_in_cooldown(self) -> bool:
        """检查是否在冷却期"""
        if self.last_scaling_time is None:
            return False
        
        elapsed = (datetime.now() - self.last_scaling_time).total_seconds()
        return elapsed < self.cooldown_seconds
    
    def _make_scaling_decision(self, metrics: LoadMetrics) -> ScalingDecision:
        """
        做出扩缩容决策
        
        Args:
            metrics: 负载指标
        
        Returns:
            扩缩容决策
        """
        utilization = metrics.get_connection_utilization()
        
        # 判断是否需要扩容
        if utilization >= self.scale_up_threshold:
            # 计算目标容量（增加50%或到达最大值）
            target = min(
                int(self.current_capacity * 1.5),
                self.max_connections
            )
            
            return ScalingDecision(
                action=ScalingAction.SCALE_UP,
                reason=f"High utilization: {utilization*100:.1f}%",
                target_connections=target,
                current_connections=self.current_capacity,
                metrics=metrics
            )
        
        # 判断是否需要缩容
        elif utilization <= self.scale_down_threshold:
            # 计算目标容量（减少30%或到达最小值）
            target = max(
                int(self.current_capacity * 0.7),
                self.min_connections
            )
            
            return ScalingDecision(
                action=ScalingAction.SCALE_DOWN,
                reason=f"Low utilization: {utilization*100:.1f}%",
                target_connections=target,
                current_connections=self.current_capacity,
                metrics=metrics
            )
        
        # 保持当前容量
        return ScalingDecision(
            action=ScalingAction.MAINTAIN,
            reason=f"Normal utilization: {utilization*100:.1f}%",
            target_connections=self.current_capacity,
            current_connections=self.current_capacity,
            metrics=metrics
        )
    
    async def _execute_scaling(self, decision: ScalingDecision):
        """
        执行扩缩容
        
        Args:
            decision: 扩缩容决策
        """
        logger.info(
            f"Executing {decision.action.value}: "
            f"{decision.current_connections} -> {decision.target_connections} "
            f"({decision.reason})"
        )
        
        # 更新容量
        old_capacity = self.current_capacity
        self.current_capacity = decision.target_connections
        
        # 记录扩缩容
        self.scaling_history.append(decision)
        if len(self.scaling_history) > 100:
            self.scaling_history = self.scaling_history[-100:]
        
        self.last_scaling_time = datetime.now()
        
        # 更新统计
        if decision.action == ScalingAction.SCALE_UP:
            self.scale_up_count += 1
        elif decision.action == ScalingAction.SCALE_DOWN:
            self.scale_down_count += 1
        
        # TODO: 实际执行扩缩容操作
        # 1. 调整连接池大小
        # 2. 如果是缩容，优雅地关闭多余连接
        # 3. 更新负载均衡配置
        
        logger.info(
            f"Scaling completed: capacity changed from {old_capacity} to {self.current_capacity}"
        )
    
    def get_current_capacity(self) -> int:
        """获取当前容量"""
        return self.current_capacity
    
    def get_recent_metrics(self, limit: int = 10) -> List[Dict[str, Any]]:
        """获取最近的指标"""
        metrics = self.recent_metrics[-limit:]
        return [m.to_dict() for m in metrics]
    
    def get_scaling_history(self, limit: int = 20) -> List[Dict[str, Any]]:
        """获取扩缩容历史"""
        history = self.scaling_history[-limit:]
        return [d.to_dict() for d in history]
    
    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        current_metrics = self.recent_metrics[-1] if self.recent_metrics else None
        
        return {
            "current_capacity": self.current_capacity,
            "min_connections": self.min_connections,
            "max_connections": self.max_connections,
            "target_utilization": round(self.target_utilization * 100, 2),
            "scale_up_count": self.scale_up_count,
            "scale_down_count": self.scale_down_count,
            "last_scaling_time": self.last_scaling_time.isoformat() if self.last_scaling_time else None,
            "current_load": current_metrics.to_dict() if current_metrics else None,
            "is_running": self.is_running
        }


# 全局实例
_autoscaler: Optional[AutoScaler] = None


def get_autoscaler(**kwargs) -> AutoScaler:
    """获取自动扩缩容器"""
    global _autoscaler
    if _autoscaler is None:
        _autoscaler = AutoScaler(**kwargs)
    return _autoscaler

